import {Annotation, END, MemorySaver, START, StateGraph} from "@langchain/langgraph";
import {tool} from "@langchain/core/tools";
import {z} from 'zod'
import {ToolNode} from "@langchain/langgraph/prebuilt";
import {AIMessage, HumanMessage} from "@langchain/core/messages";
import {qwen} from "../models";
import {saveImages} from "../tool";


const AgentState = Annotation.Root({
    messages: Annotation<string>({
        reducer: (x, y)=> x.concat(y)
    })
})

const memory = new MemorySaver()


const searchTool = tool((_):string=> {
    return "It's sunny in San Francisco, but you better look out if you're a Gemini 😈."
}, {
    name: 'search',
    description: 'Call to surf the web.',
    schema: z.object({
        query: z.string()
    })
})

const tools = [searchTool]
const toolNode = new ToolNode<typeof AgentState.State>(tools)
const boundModel = qwen.bindTools(tools)

function shouldContinue(state: typeof AgentState.State): 'action' | typeof END {
    const lastMessage = state.messages[state.messages.length -1];

    if(lastMessage && !(lastMessage as AIMessage).tool_calls?.length) {
        return END;
    }
    return 'action'
}

async function callModel(state: typeof AgentState.State) {
    const response = await qwen.invoke(state.messages);
    return {
        messages: [response]
    }
}

const workflow = new StateGraph(AgentState)
.addNode("agent", callModel)
.addNode("action", toolNode)
.addConditionalEdges("agent", shouldContinue)
.addEdge("action", "agent")
.addEdge(START, "agent")


const app = workflow.compile({
    checkpointer: memory
})

// await saveImages(app, 'memory.png')

const config = {
    configurable: {
        thread_id: "2",
        streamMode: 'values'
    }
}

const inputMessage = new HumanMessage("hi! I'm bob")

for await (const event of await app.stream({
    messages: [inputMessage]
}, config)) {
    const recentMsg = event.agent.messages[event.agent.messages.length - 1];
    console.log(`================================ ${recentMsg._getType()} Message (1) =================================`)
    console.log(recentMsg.content);
}

console.log("\n\n================================= END =================================\n\n")

const inputMessage2 = new HumanMessage("what's my name?");
for await (const event of await app.stream({
    messages: [inputMessage2]
}, config)) {
    const recentMsg = event.agent.messages[event.agent.messages.length - 1];
    console.log(`================================ ${recentMsg._getType()} Message (2) =================================`)
    console.log(recentMsg.content);
}